#include "get_arg.h"
#include <boost/filesystem/path.hpp>
#include <boost/filesystem/operations.hpp>
#include <boost/foreach.hpp>
#include <iostream>
#include <stdlib.h>
#include <sstream>
#include <fstream>
#include <sparse_matrix.h>
#include <sparse_matrix_io.h>
#include <portable_timer.h>
#include <boost/lexical_cast.hpp>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <tree.h>


using namespace std;
using namespace bcpp_matrix;
using boost::lexical_cast;
using std::ostringstream;
using std::istringstream;
using std::ofstream;
using std::ifstream;

unsigned get_file_size(const string& name)
{
	struct stat file_stat;
	if(stat(name.c_str(), &file_stat))
		return 0;
	return file_stat.st_size;

}

struct t_iter_index
{
	unsigned index;
	vector<string>::const_iterator iter;
	t_iter_index(unsigned i, const vector<string>::const_iterator& it)
		:index(i), iter(it){}
	bool operator<( const t_iter_index& other)
	{
		return *(other.iter) < *iter;
	}
	operator unsigned () const {return index;}

};


struct t_output_message_with_timer
{
	t_portable_timer timer;
	string end_message;
	unsigned size;
	unsigned pos;
	void set_size(unsigned i) {size = i;}
	t_output_message_with_timer(const string& message, bool delay = false)
		: size(0), pos(0)
	{
		if (delay)
		{
			end_message = message;
		}
		else
		{
			pos = message.length() + 8;
			cerr << "\t" << message;
		}
	}
	~t_output_message_with_timer()
	{
		if (end_message.length())
		{
			pos += end_message.length() + 8;
			cerr << "\t" << end_message;
		}

		if (size)
		{
			ostringstream str;
			if (size > 1024*1024*10)
				str << " (" << size / 1024 / 1024 << " MB) ";
			else if (size > 1024 * 10)
				str << " (" << size / 1024 << " kB) ";
			else
				str << "(" << size << " bytes) ";
			pos += str.str().length();
			cerr << str.str();
		}
		if (pos < 70)
			cerr << string(70 - pos, ' ');
		cerr << timer.pretty_reset() << "\n";
	}
};


void get_num_values(vector<unsigned>& counts, unsigned desired_total)
{
	unsigned curr_total = 0;
	unsigned MATRIX_SZ = counts.size();
	desired_total  = std::max(desired_total,  MATRIX_SZ);
	for (unsigned i = 0; i < MATRIX_SZ; ++i)
	{
		counts[i] = i;
		curr_total += i;
	}
	while (curr_total > desired_total)
	{
		unsigned col = rand() %  MATRIX_SZ;
		// each row has at least 1 element
		if (counts[col] <= 1)
			continue;
		unsigned possible = std::min(counts[col] - 1, curr_total - desired_total);
		unsigned remove = rand() % possible + 1;
		curr_total -=remove;
		counts[col] -= remove;
	}
}
void randomize_counts(vector<unsigned>& counts)
{
	unsigned sz = counts.size();
	for (unsigned i = 0U; i < sz; ++i)
	{
		unsigned i = rand() % sz;
		unsigned j = (i + rand()) % sz;
		std::swap(counts[i], counts[j]);
	}
	
}

struct t_val
{
	unsigned	i; 
	unsigned	j; 
	float		f;
	bool operator < (const t_val& o) const
	{
		if (i == o.i)
			return j < o.j;
		else
			return i < o.i;
	}
	t_val(unsigned ii, unsigned jj, float ff): i(ii), j(jj), f(ff){	}
};



void test_tree_read(std::istream& f_temp, const std::string& tree_str)
{
	bcpp_tree::t_tree input_tree;
	input_tree.set_sort_by_names(true);
	std::string input_tree_str;
	std::string tree_name;
	unsigned cnt_lines = 0;
	bool success = read_tree_from_stream(f_temp, input_tree, cnt_lines);
	assert(success);
	input_tree_str = ">" + input_tree.get_name() + "\n" + input_tree.to_nh_str() + "\n";
	if (input_tree_str != tree_str)
	{
		cerr << "\n>>original tree_str = [" << tree_str << "]\n";
		cerr << "\n>>read tree_str = [" << input_tree_str << "]\n";
	}
	assert(input_tree_str == tree_str);

}


int main (int argc, char *argv[])
{

	std::string tree_str (">test_matrix\n((a:0.50000,b:0.50000):0.50000,c:0.50000);\n");
	t_program_args args;
	if (!process_command_line_options(argc, argv, args))
		return 1;

	const unsigned MATRIX_SZ = args.size;
	const float DEFAULT_VAL = args.default_val;
	const unsigned cnt_values = MATRIX_SZ * (MATRIX_SZ - 1) / 2;
	const unsigned PROPORTION = min(cnt_values, unsigned(cnt_values * args.fill));
	cerr << "\tTo fill " << PROPORTION << " diagonal elements(";
	cerr << "\t" << PROPORTION * 2 + MATRIX_SZ << " elements)\n";


	//
	// Allocate elements
	//
	t_sparse_matrix_f matrix1;
	{
		cerr << "\tAllocate " << MATRIX_SZ << " elements\n";
		matrix1.allocate(MATRIX_SZ);
	}

	//
	// Assign names
	//
	{
		t_output_message_with_timer message("Assign names");
		for (unsigned i = 0; i < MATRIX_SZ; ++i)
			matrix1.names[i] = "A_" + lexical_cast<string>(i);
		assert(matrix1.size() == MATRIX_SZ);
	}
	matrix1.matrix_name = "test1";

	deque<t_val> values;

	//
	// Randomly assign values
	//
	{
		t_output_message_with_timer message("Randomly assign values\n");

		// diagonals are zeros to be compatible with triangle format
		for (unsigned ii = 0; ii < MATRIX_SZ; ++ii)
			values.push_back(t_val(ii, ii, 0.0));

		// get number of values per row (randomly)
		vector<unsigned> counts(MATRIX_SZ, 0);
		get_num_values(counts, PROPORTION);
		vector<unsigned> indices;
		for (unsigned i = 0; i < MATRIX_SZ; ++i)
			indices.push_back(i);
		std::cerr << "\tNumber of values per row decided\n";

		t_progress_indicator dots(args.VERBOSE(), (std::max(PROPORTION,  MATRIX_SZ)) / 100 + 1);
		dots.set_limit(std::max(PROPORTION,  MATRIX_SZ));
		dots.use_timer();
		unsigned cnt_filled	= 0;
		for (unsigned row = 0; row < MATRIX_SZ; ++row)
		{
			unsigned count_values_in_row = counts[row];
			randomize_counts(indices);
			vector<unsigned> sorted_indices(count_values_in_row);
			std::copy(indices.begin(),  indices.begin() + count_values_in_row,  sorted_indices.begin());
			std::sort(sorted_indices.begin(),  sorted_indices.end());
			for (unsigned i = 0; i < count_values_in_row; ++i)
			{
				unsigned col = sorted_indices[i];
				float val = rand() * 10.0 /  RAND_MAX;
				values.push_back(t_val(row, col, val));
				values.push_back(t_val(col, row, val));

				++dots;
				++cnt_filled;
			}
		}
		dots.finish();
		cerr << "\t\t" << cnt_filled << " diagonal values filled.\n";
	}
	{
		std::cerr << "\tSorting paired coordinates before entering into matrix\n";
		sort(values.begin(),  values.end());
		t_progress_indicator dots(args.VERBOSE(), (std::max(PROPORTION,  MATRIX_SZ)) / 100 + 1);
		dots.set_limit((std::max(PROPORTION,  MATRIX_SZ)));
		dots.use_timer();
		BOOST_FOREACH(t_val& val, values) 
		{
			++dots;
			matrix1.matrix(val.i, val.j) = val.f;
		}
	}


	cerr << "\t\t" << matrix1.get_count_filled() << " cells filled.\n";
	cerr << "\t\t" << matrix1.get_count_missing() << " missing out of ";
	cerr << matrix1.get_count_pairs() << " expected\n";


	cerr << "\t\t" << matrix1.get_count_missing() << " missing out of ";
	cerr << matrix1.get_count_pairs() << " expected\n";

	//
	// Round values so that comparisons will work
	//
	{
		t_output_message_with_timer message("Round up values by writing out then reading back");
		stringstream strm;
		write_matrix(strm, matrix1, eMIO_PAIRS, DEFAULT_VAL);
		read_matrix(strm, matrix1, eMIO_PAIRS, DEFAULT_VAL);
	}

	//
	// Sort names
	//
	{
		
		t_output_message_with_timer message("Sort names");
		matrix1.sort_by_name_alpha();
	}

	//
	// pair-wise format
	//
	{
		cerr << "\n";
		t_output_message_with_timer message("Matrix in pairwise format", true);
		string data;
		t_sparse_matrix_f matrix2;
		{
			ofstream f_temp(args.temp_file.c_str());
			if (!f_temp)
			{
				cerr << "Error:\n\tCould not open " << args.temp_file << "\n";
				return 1;
			}
			t_output_message_with_timer message("Write matrix in pair-wise format");
			write_matrix(f_temp, matrix1, eMIO_PAIRS, DEFAULT_VAL);
			f_temp << tree_str;
			f_temp.close();
			message.set_size(get_file_size(args.temp_file));
		}
		{
			ifstream f_temp(args.temp_file.c_str());
			if (!f_temp)
			{
				cerr << "Error:\n\tCould not open " << args.temp_file << "\n";
				return 1;
			}
			t_output_message_with_timer message("Read matrix in pair-wise format");
			read_matrix(f_temp, matrix2, eMIO_PAIRS, DEFAULT_VAL);
		}
		{
			t_output_message_with_timer message("Sort matrix for comparison");
			matrix2.sort_by_name_alpha();
		}
		assert(matrix2 == matrix1);
	}



//	write_matrix_to_square(cout, matrix1, true, false);

	//
	// binary format
	//
	{
		cerr << "\n";
		t_output_message_with_timer message("Matrix in binary format", true);
		string data;
		t_sparse_matrix_f matrix2;
		{
			ofstream f_temp(args.temp_file.c_str());
			if (!f_temp)
			{
				cerr << "Error:\n\tCould not open " << args.temp_file << "\n";
				return 1;
			}
			t_output_message_with_timer message("Write matrix in binary format");
			write_matrix(f_temp, matrix1, eMIO_BINARY, DEFAULT_VAL);
			//write_matrix(f_temp, matrix1, eMIO_SQUARE, false);
			f_temp << tree_str;
			f_temp.close();
			message.set_size(get_file_size(args.temp_file));
		}
		{
			ifstream f_temp(args.temp_file.c_str());
			if (!f_temp)
			{
				cerr << "Error:\n\tCould not open " << args.temp_file << "\n";
				return 1;
			}
			t_output_message_with_timer message("Read matrix in binary format");
			read_matrix(f_temp, matrix2, eMIO_BINARY, DEFAULT_VAL);
			test_tree_read(f_temp, tree_str);
		}
		assert(matrix2.names == matrix1.names);
		//for (unsigned i = 0; i < MATRIX_SZ; ++i)
		//{
		//	for (unsigned j = 0; j < MATRIX_SZ; ++j)
		//	{
		//		cerr << "[" << i << "," << j << "] = " << matrix1.matrix[i][j] << "\t"
		//				 << matrix2.matrix[i][j]
		//				<<	(matrix1.is_missing(i, j) ? " missing1": " present1") << "\t"
		//				<<	(matrix2.is_missing(i, j) ? " missing2": " present2") << "\n";
		//	}
		//}
		cerr << "matrix1 missing count = " << matrix1.get_count_missing() << "\n";
		cerr << "matrix1 missing count = " << matrix2.get_count_missing() << "\n";
		assert(matrix2 == matrix1);
	}

	//
	// Lower triangular format
	//
	{
		cerr << "\n";
		t_output_message_with_timer message("Matrix in lower triangular format", true);
		string data;
		t_sparse_matrix_f matrix2;
		{
			ofstream f_temp(args.temp_file.c_str());
			if (!f_temp)
			{
				cerr << "Error:\n\tCould not open " << args.temp_file << "\n";
				return 1;
			}
			t_output_message_with_timer message("Write matrix in lower triangular format");
			write_matrix(f_temp, matrix1, eMIO_LOWER, DEFAULT_VAL);
			f_temp << tree_str;
			f_temp.close();
			message.set_size(get_file_size(args.temp_file));
		}
		{
			ifstream f_temp(args.temp_file.c_str());
			if (!f_temp)
			{
				cerr << "Error:\n\tCould not open " << args.temp_file << "\n";
				return 1;
			}
			t_output_message_with_timer message("Read matrix in lower triangular format");
			read_matrix(f_temp, matrix2, eMIO_LOWER, DEFAULT_VAL);
			test_tree_read(f_temp, tree_str);
		}
		{
			ofstream f_temp(args.temp_file.c_str(), std::ios::app | std::ios::ate);
			write_matrix(f_temp, matrix1, eMIO_SQUARE, DEFAULT_VAL);
			write_matrix(f_temp, matrix2, eMIO_SQUARE, DEFAULT_VAL);
		}
		assert(matrix2 == matrix1);
	}

	//
	// Square format
	//
	{
		cerr << "\n";
		t_output_message_with_timer message("Matrix in square format", true);
		string data;
		t_sparse_matrix_f matrix2;
		{
			ofstream f_temp(args.temp_file.c_str());
			if (!f_temp)
			{
				cerr << "Error:\n\tCould not open " << args.temp_file << "\n";
				return 1;
			}
			t_output_message_with_timer message("Write matrix in square format");
			write_matrix(f_temp, matrix1, eMIO_SQUARE, DEFAULT_VAL);
			f_temp << tree_str;
			f_temp.close();
			message.set_size(get_file_size(args.temp_file));
		}
		{
			ifstream f_temp(args.temp_file.c_str());
			if (!f_temp)
			{
				cerr << "Error:\n\tCould not open " << args.temp_file << "\n";
				return 1;
			}
			t_output_message_with_timer message("Read matrix in square format");
			read_matrix(f_temp, matrix2, eMIO_SQUARE, DEFAULT_VAL);
			test_tree_read(f_temp, tree_str);
		}
		assert(matrix2 == matrix1);
	}

	boost::filesystem::remove(boost::filesystem::path(args.temp_file));;

	cerr << "\n";

	// Multiple formats
	//
	{
		cerr << "\n";
		t_output_message_with_timer message("Matrix in multiple formats", true);

		string data;
		t_sparse_matrix_f matrix2;
		{
			ofstream f_temp(args.temp_file.c_str());
			if (!f_temp)
			{
				cerr << "Error:\n\tCould not open " << args.temp_file << "\n";
				return 1;
			}
			{
				matrix1.sort_by_name_alpha();
				t_output_message_with_timer message("Write matrix in multiple formats");
				cerr << "!1\tSquare format\n";
				write_matrix(f_temp, matrix1, eMIO_SQUARE, DEFAULT_VAL);
				f_temp << tree_str;
				cerr << "!1\tBinary format\n";
				write_matrix(f_temp, matrix1, eMIO_BINARY, DEFAULT_VAL);
				f_temp << tree_str;
				cerr << "!1\tLower format\n";
				write_matrix(f_temp, matrix1, eMIO_LOWER, DEFAULT_VAL);
				f_temp << tree_str;
				cerr << "!1\tPairs format\n";
				write_matrix(f_temp, matrix1, eMIO_PAIRS, DEFAULT_VAL);
				f_temp << tree_str;
				cerr << "!2\tSquare format\n";
				write_matrix(f_temp, matrix1, eMIO_SQUARE, DEFAULT_VAL);
				f_temp << tree_str;
				cerr << "!2\tBinary format\n";
				write_matrix(f_temp, matrix1, eMIO_BINARY, DEFAULT_VAL);
				f_temp << tree_str;
				cerr << "!2\tLower format\n";
				write_matrix(f_temp, matrix1, eMIO_LOWER, DEFAULT_VAL);
				f_temp << tree_str;
				cerr << "!2\tPairs format\n";
				write_matrix(f_temp, matrix1, eMIO_PAIRS, DEFAULT_VAL);
				f_temp << tree_str;
				cerr << "!3\tSquare format\n";
				write_matrix(f_temp, matrix1, eMIO_SQUARE, DEFAULT_VAL);
				f_temp << tree_str;
			}
			f_temp.close();
			message.set_size(get_file_size(args.temp_file));
		}
		{
			ifstream f_temp(args.temp_file.c_str());
			if (!f_temp)
			{
				cerr << "Error:\n\tCould not open " << args.temp_file << "\n";
				return 1;
			}
			{
				t_output_message_with_timer message("Read matrix in multiple formats");
				std::ostringstream ignore;
				t_progress_indicator dots(ignore, 1);
				unsigned cnt_lines = 0;

				cerr << ">1\tSquare format\n";
				f_temp >> t_sparse_matrix_istrm_wrapper_f(matrix2, eMIO_SQUARE,  cnt_lines,  dots);
				//read_matrix(f_temp, matrix2, eMIO_SQUARE, ignore_cnt_duplicates, DEFAULT_VAL);
				assert(matrix2 == matrix1);
				test_tree_read(f_temp, tree_str);

				cerr << ">1\tBinary format\n";
				read_matrix(f_temp, matrix2, eMIO_BINARY, DEFAULT_VAL);
				assert(matrix2 == matrix1);
				test_tree_read(f_temp, tree_str);

				cerr << ">1\tLower format\n";
				read_matrix(f_temp, matrix2, eMIO_LOWER, DEFAULT_VAL);
				assert(matrix2 == matrix1);
				test_tree_read(f_temp, tree_str);

				cerr << ">1\tPairs format\n";
				read_matrix(f_temp, matrix2, eMIO_PAIRS, DEFAULT_VAL);
				matrix2.sort_by_name_alpha();
				if (!(matrix2 == matrix1))
				{
					std::cerr << "first\n";
					write_matrix(std::cerr, matrix1, eMIO_PAIRS, DEFAULT_VAL);
					std::cerr << "second\n";
					write_matrix(std::cerr, matrix2, eMIO_PAIRS, DEFAULT_VAL);
				}
				assert(matrix2 == matrix1);
				test_tree_read(f_temp, tree_str);

				cerr << ">2\tSquare format\n";
				read_matrix(f_temp, matrix2, eMIO_SQUARE, DEFAULT_VAL);
				assert(matrix2 == matrix1);
				test_tree_read(f_temp, tree_str);

				cerr << ">2\tBinary format\n";
				read_matrix(f_temp, matrix2, eMIO_BINARY, DEFAULT_VAL);
				assert(matrix2 == matrix1);
				test_tree_read(f_temp, tree_str);

				cerr << ">2\tLower format\n";
				read_matrix(f_temp, matrix2, eMIO_LOWER, DEFAULT_VAL);
				assert(matrix2 == matrix1);
				test_tree_read(f_temp, tree_str);

				cerr << ">2\tPairs format\n";
				read_matrix(f_temp, matrix2, eMIO_PAIRS, DEFAULT_VAL);
				matrix2.sort_by_name_alpha();
				assert(matrix2 == matrix1);
				test_tree_read(f_temp, tree_str);

				cerr << ">3\tSquare format\n";
				read_matrix(f_temp, matrix2, eMIO_SQUARE, DEFAULT_VAL);
				assert(matrix2 == matrix1);
				test_tree_read(f_temp, tree_str);
			}
		}
	}




	cerr << "\n";
	cerr << string(10, '!') << "\nEverything worked!!\n";

	return 0;
}


